import pickle
from dataclasses import dataclass
from typing import Any, Callable, Generic, TypeVar, Union

import torch
from numpy.typing import ArrayLike
from zkl_aiutils_training import ProcessingTaskPlugin, Resumable
from zkl_ptutils_metrics import MeanMetric, Metric

from zkl_ptutils_training.utils.attributes import get_entry_or_attribute
from .metrics import MetricScope, MetricsProvider
from .resumable_fs import FsPauseArgs, FsResumeArgs, FsResumeFromCheckpointArgs

AnyValue = TypeVar("AnyValue")
AnyInput = TypeVar('AnyInput')
AnyOutput = TypeVar('AnyOutput')
AnyCreateArgs = TypeVar('AnyCreateArgs')
MetricPatchValue = Union[torch.Tensor, ArrayLike, Any]


@dataclass
class TypedMetricPatch(Generic[AnyValue]):
    value: AnyValue
    typ: Callable[[], Metric] | type[Metric] | Metric = MeanMetric


@dataclass
class ScopedTypedMetricPatch(TypedMetricPatch[AnyValue]):
    scope: MetricScope = 'epoch'


class StepMetric(Metric[int, int]):
    def __init__(self, value: int = 0):
        self._value = value

    def compute(self) -> int:
        return self._value

    def update(self, patch: int = 1):
        self._value += patch

    def reset(self, value: int = 0):
        self._value = value


class MetricsRecording(ProcessingTaskPlugin[AnyInput, AnyOutput], MetricsProvider, Resumable):
    def __init__(self):
        self._scoped_metrics: dict[str, tuple[Metric, MetricScope]] = {}

    # resume & pause

    def on_resume(self, args: FsResumeArgs):
        if isinstance(args, FsResumeFromCheckpointArgs):
            with args.checkpoint_fs.open('metrics.pkl', 'rb') as fp:
                self._scoped_metrics.update(pickle.load(fp))

    def on_pause(self, args: FsPauseArgs):
        with args.checkpoint_fs.open('metrics.pkl', 'wb') as fp:
            pickle.dump(self._scoped_metrics, fp)

    # step

    def on_after_process(self, output: AnyOutput) -> AnyOutput:
        output = super().on_after_process(output)
        metrics = get_entry_or_attribute(output, 'metrics', {})
        if 'step' not in metrics:
            metrics['step'] = ScopedTypedMetricPatch(1, StepMetric, 'run')
        self._update_metrics(metrics)
        return output

    # metrics

    def _get_scoped_metrics(self) -> dict[str, tuple[Metric, MetricScope]]:
        return self._scoped_metrics

    def _update_metrics(self, metrics: dict[str, Any | TypedMetricPatch | ScopedTypedMetricPatch]):
        for name, patch in metrics.items():
            self._update_metric(name, patch)

    def _update_metric(self, name: str, patch: Any | TypedMetricPatch | ScopedTypedMetricPatch):
        scope: MetricScope
        scope = patch.scope if isinstance(patch, ScopedTypedMetricPatch) else 'epoch'
        typ = patch.typ if isinstance(patch, TypedMetricPatch) else MeanMetric
        value = patch.value if isinstance(patch, TypedMetricPatch) else patch

        scoped_metric = self._scoped_metrics.get(name)
        if scoped_metric is None:
            metric = typ if isinstance(typ, Metric) else typ()
            self._scoped_metrics[name] = (metric, scope)
        else:
            metric, _ = scoped_metric

        metric.update(value)
